{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct Shooting vs Direct Transcription"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import pydrake\n",
    "    import underactuated\n",
    "except ImportError:\n",
    "    !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "    from jupyter_setup import setup_underactuated\n",
    "    setup_underactuated()\n",
    "\n",
    "# Setup matplotlib backend (to notebook, if possible, or inline).\n",
    "from underactuated.jupyter import setup_matplotlib_backend\n",
    "plt_is_interactive = setup_matplotlib_backend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# python libraries\n",
    "import sys\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "from scipy.sparse.linalg import spsolve\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# pydrake imports\n",
    "from pydrake.all import DiscreteTimeLinearQuadraticRegulator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "\n",
    "In this notebook we compare two approaches to transcribe optimal control problems ([see here for more details](http://underactuated.mit.edu/trajopt.html#computational_considerations)):\n",
    "- **Direct transcription,** where in the optimal control problem both the inputs $\\mathbf{u}[n]$ and the states $\\mathbf{x}[n]$ are decision variables, and the dynamic equations $\\mathbf{x}[n+1] = \\mathbf{f}(\\mathbf{x}[n], \\mathbf{u}[n])$ are kept as constraints.\n",
    "- **Direct shooting,** where the states $\\mathbf{x}[n]$ are expressed as functions of the initial state $\\mathbf{x}[0]$ and the input sequence $(\\mathbf{u}[0], \\ldots, \\mathbf{u}[n-1])$, and eliminated from the optimization problem.\n",
    "\n",
    "As discussed in the text, in some cases, direct transcription can be numerically more robust and also more efficient, even if direct shooting yields more compact optimization problem (less variables and constraints).\n",
    "In this notebook we analyze the typical example in which this performance gap appears: optimal control of unstable systems with \"long\" time horizon.\n",
    "\n",
    "To make the analysis as simple as possible, we will consider the finite-horizon LQR problem:\n",
    "\\begin{align}\n",
    "\\min & \\quad J = \\mathbf{x}^T[N] \\mathbf{Q}_N \\mathbf{x}[N]  + \\sum_{n=0}^{N-1} (\\mathbf{x}^T[n] \\mathbf{Q} \\mathbf{x}[n] + \\mathbf{u}^T[n] \\mathbf{R} \\mathbf{u}[n]) \\\\\n",
    "\\text{subject to} & \\quad \\mathbf{x}[0] = \\mathbf{x}_0, \\\\\n",
    "& \\quad \\mathbf{x}[n+1] = \\mathbf{A} \\mathbf{x}[n] + \\mathbf{B} \\mathbf{u}[n]. \\\\\n",
    "\\end{align}\n",
    "where $\\mathbf{Q}_N \\succeq 0$, $\\mathbf{Q} \\succeq 0$, and $\\mathbf{R} \\succ 0$.\n",
    "Of course, numerical optimization is too big of a hammer for this problem: we can solve LQR in almost zero time by working on the Riccati equations!\n",
    "However, this problem setup is rich enough to exhibit all the characteristic failures that our direct methods might encounter when working with complex robots such as a humanoid.\n",
    "\n",
    "In this notebook, you'll be asked to work on four pieces of code:\n",
    "- analyze the conditioning of the linear system of equations we get when using direct shooting,\n",
    "- implement part of the direct transcription method,\n",
    "- verify that the cost to go from direct trascription approaches the LQR cost to go as the time horizon grows,\n",
    "- analyze the conditioning of direct transcription,\n",
    "- implement the dynamic programming solution (a.k.a. Riccati recursion) for this problem.\n",
    "\n",
    "Don't worry if you aren't quite sure about the meaning of \"conditioning,\" you'll find all the details below!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmark System\n",
    "\n",
    "Just to fix the ideas and plot some results, here we define a dynamical system which we'll use to test the functions we'll write in the notebook.\n",
    "In particular, we use the discretized and linearized version of the inverted pendulum dynamics.\n",
    "As we will see, unstable systems clearly highlight the advantages of direct transcription methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# time discretization step\n",
    "h = .5\n",
    "\n",
    "# linearized continuous time dynamics\n",
    "# (all physical parameters of the pendulum are set to one)\n",
    "Act = np.array([[0, 1],[1, 0]])\n",
    "Bct = np.array([[0],[1]])\n",
    "\n",
    "# discretize using explicit Euler\n",
    "# x[n+1] = x[n] + h (Act x[n] + Bct u[n])\n",
    "#        = (I + h Act) x[n] + h Bct u[n]\n",
    "A = np.eye(Act.shape[0]) + h * Act\n",
    "B = h * Bct\n",
    "\n",
    "# objective function\n",
    "Q = np.array([[1, 0], [0, 0]])\n",
    "R = np.array([[1]])\n",
    "QN = np.array([[3, 0], [0, 2]])\n",
    "\n",
    "# initial state\n",
    "x0 = np.array([1, 1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an asymptotic comparison for the different optimization problems we'll solve, we use the \"closed form\" solution of the infinite-horizon discrete-time Riccati equations.\n",
    "Recall that our LQR problem above has finite time horizon $N$, hence we expect its solution to converge to the infinite-horizon only for long-enough horizons $N$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimal feedback gain and cost-to-go matrix\n",
    "K, S = DiscreteTimeLinearQuadraticRegulator(A, B, Q, R)\n",
    "\n",
    "# cost to go from state x0\n",
    "J_star_inf = x0.dot(S).dot(x0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Direct shooting\n",
    "\n",
    "We start with direct shooting.\n",
    "[In the lecture notes](http://underactuated.mit.edu/trajopt.html#direct_shooting), we have seen that the state $\\mathbf{x}[n]$ can be expressed as a function of the initial state $\\mathbf{x}[0]$ and the input sequence $\\mathbf{u}[0], \\ldots, \\mathbf{u}[n-1]$.\n",
    "We derived the formula\n",
    "$$\\mathbf{x}[n] = \\mathbf{A}^n \\mathbf{x}_0 + \\sum_{k=0}^{n-1} \\mathbf{A}^{n-k-1} \\mathbf{B} \\mathbf{u}[k].$$\n",
    "\n",
    "Defining the vectors \n",
    "$$\n",
    "\\mathbb{x} =\n",
    "\\begin{bmatrix}\n",
    "\\mathbf{x}[1] \\\\\n",
    "\\vdots \\\\\n",
    "\\mathbf{x}[N]\n",
    "\\end{bmatrix},\n",
    "\\quad\n",
    "\\mathbb{u} =\n",
    "\\begin{bmatrix}\n",
    "\\mathbf{u}[0] \\\\\n",
    "\\vdots \\\\\n",
    "\\mathbf{u}[N-1]\n",
    "\\end{bmatrix},\n",
    "$$\n",
    "the formula above can be written in matrix form as\n",
    "$$\n",
    "\\mathbb{x} =\n",
    "\\begin{bmatrix}\n",
    "\\mathbf{A} \\\\\n",
    "\\vdots \\\\\n",
    "\\mathbf{A}^N\n",
    "\\end{bmatrix}\n",
    "\\mathbf{x}_0 +\n",
    "\\begin{bmatrix}\n",
    "\\mathbf{B} & 0 & \\cdots & 0 \\\\\n",
    "\\mathbf{A B} & \\mathbf{B} & \\cdots & 0 \\\\\n",
    "\\vdots & \\vdots & \\ddots  & \\vdots \\\\\n",
    "\\mathbf{A}^{N-1} \\mathbf{B} & \\mathbf{A}^{N-2} \\mathbf{B} & \\cdots & \\mathbf{B} \\\\\n",
    "\\end{bmatrix}\n",
    "\\mathbb{u}.\n",
    "$$\n",
    "Let us call these two matrices $\\mathbb{A}$ and $\\mathbb{B}$ (in this notebook we'll use $\\mathbb{blackboard \\ bold}$ letters for block matrices), so that $\\mathbb{x} = \\mathbb{A} \\mathbf{x}_0 + \\mathbb{B} \\mathbb{u}$.\n",
    "\n",
    "In the following cell we provide two efficient functions to construct $\\mathbb{A}$ and $\\mathbb{B}$.\n",
    "In the code we denote, e.g., the matrix $\\mathbb{A}$ as `Ab`, where `b` stands for block matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# efficient computation of the matrix Ab\n",
    "def get_Ab(A, N):\n",
    "    \n",
    "    # Ab is defined only for N >= 1\n",
    "    assert N >= 1\n",
    "    \n",
    "    # number of states\n",
    "    p = A.shape[0]\n",
    "    \n",
    "    # initialize matrix with all zeros\n",
    "    Ab = np.zeros((p * N, p))\n",
    "    \n",
    "    # put A matrix in top block\n",
    "    rows = slice(0, p) # equivalent to :p\n",
    "    Ab[rows, :] = A\n",
    "    \n",
    "    # loop over the block-rows\n",
    "    for n in range(1, N):\n",
    "        \n",
    "        # move the row indices one block down\n",
    "        prev_rows = rows\n",
    "        rows = slice(prev_rows.start + p, prev_rows.stop + p) # equivalent to n*p:(n+1)*p\n",
    "        \n",
    "        # premultiply previous block by A\n",
    "        # place the result in this block\n",
    "        Ab[rows, :] = A.dot(Ab[prev_rows, :])\n",
    "        \n",
    "    return Ab\n",
    "\n",
    "# efficient computation of the matrix Bb\n",
    "def get_Bb(A, B, N):\n",
    "    \n",
    "    # Bb is defined only for N >= 1\n",
    "    assert N >= 1\n",
    "    \n",
    "    # number of states and inputs\n",
    "    p, q = B.shape\n",
    "    \n",
    "    # initialize matrix with all zeros\n",
    "    Bb = np.zeros((p * N, q * N))\n",
    "    \n",
    "    # put B matrix in top-left block\n",
    "    rows = slice(0, p) # equivalent to :p\n",
    "    Bb[rows, :q] = B\n",
    "    \n",
    "    # loop over the block-rows\n",
    "    for n in range(1, N):\n",
    "        \n",
    "        # move the row indices one block down\n",
    "        prev_rows = rows\n",
    "        rows = slice(prev_rows.start + p, prev_rows.stop + p) # equivalent to n*p:(n+1)*p\n",
    "        \n",
    "        # premultiply first block from previous block-row by A\n",
    "        # place the result as first block of this block-row\n",
    "        Bb[rows, :q] = A.dot(Bb[prev_rows, :q])\n",
    "        \n",
    "        # shift the rest of previous block-row one block forward and\n",
    "        # place it in the current block-row\n",
    "        Bb[rows, q:(n+1)*q] = Bb[prev_rows, :n*q]\n",
    "        \n",
    "    return Bb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us write the LQR objective function in block-matrix form:\n",
    "$$\n",
    "J =\n",
    "\\mathbf{x}_0^T \\mathbf{Q} \\mathbf{x}_0 +\n",
    "\\mathbb{x}^T\n",
    "\\begin{bmatrix}\n",
    "\\mathbf{Q} & \\cdots & 0 & 0 \\\\\n",
    "\\vdots & \\ddots & \\vdots & \\vdots \\\\\n",
    "0 & \\cdots  & \\mathbf{Q} & 0 \\\\\n",
    "0 & \\cdots  & 0 & \\mathbf{Q}_N \\\\\n",
    "\\end{bmatrix}\n",
    "\\mathbb{x} +\n",
    "\\mathbb{u}^T\n",
    "\\begin{bmatrix}\n",
    "\\mathbf{R} & \\cdots & 0 \\\\\n",
    "\\vdots & \\ddots  & \\vdots \\\\\n",
    "0 & \\cdots  & \\mathbf{R}\\\\\n",
    "\\end{bmatrix}\n",
    "\\mathbb{u}.\n",
    "$$\n",
    "We denote the two block-diagonal matrices in the latter equation as $\\mathbb{Q}$ and $\\mathbb{R}$, so that $J = \\mathbf{x}_0^T \\mathbf{Q} \\mathbf{x}_0 + \\mathbb{x}^T \\mathbb{Q} \\mathbb{x} + \\mathbb{u}^T \\mathbb{R} \\mathbb{u}$.\n",
    "\n",
    "We now use the relation $\\mathbb{x} = \\mathbb{A} \\mathbf{x}_0 + \\mathbb{B} \\mathbb{u}$.\n",
    "After a simple manipulation, we get\n",
    "$$\n",
    "J = \\mathbf{x}_0^T (\\mathbf{Q} + \\mathbb{A}^T \\mathbb{Q} \\mathbb{A}) \\mathbf{x}_0 +\n",
    "2 \\mathbf{x}_0^T \\mathbb{A}^T \\mathbb{Q} \\mathbb{B} \\mathbb{u} +\n",
    "\\mathbb{u}^T (\\mathbb{R} + \\mathbb{B}^T \\mathbb{Q} \\mathbb{B}) \\mathbb{u}.\n",
    "$$\n",
    "Where we used the fact that the matrices $\\mathbf{Q}$, $\\mathbf{R}$, and $\\mathbf{Q}_N$ are symmetric (hence, so are $\\mathbb{Q}$ and $\\mathbb{R}$).\n",
    "\n",
    "The finite-time LQR problem has now become an unconstrained quadratic minimization problem.\n",
    "To find the optimal control sequence $\\mathbb{u}^*$, we just set the derivative of $J$ with respect to $\\mathbb{u}$ to zero, and solve the resulting linear system.\n",
    "We get\n",
    "$$\n",
    "2 \\mathbb{B}^T \\mathbb{Q} \\mathbb{A} \\mathbf{x}_0 +\n",
    "2 (\\mathbb{R} + \\mathbb{B}^T \\mathbb{Q} \\mathbb{B}) \\mathbb{u}\n",
    "= 0.\n",
    "$$\n",
    "Let use define $\\mathbb{H} = \\mathbb{R} + \\mathbb{B}^T \\mathbb{Q} \\mathbb{B}$.\n",
    "Note that this matrix is positive definite (hence invertible), since $\\mathbb{R} \\succ 0$ and $\\mathbb{B}^T \\mathbb{Q} \\mathbb{B} \\succeq 0$.\n",
    "The optimal control sequence is then\n",
    "$$ \\mathbb{u}^* = - \\mathbb{H}^{-1} \\mathbb{B}^T \\mathbb{Q} \\mathbb{A} \\mathbf{x}_0.$$\n",
    "The optimal cost $J^*$ can be derived by plugging the controls $\\mathbb{u}^*$ in the expression of $J$.\n",
    "\n",
    "In the following cell we implemented these steps for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that given the system state x0 (together with:\n",
    "# the dynamical system matrices A, B; the objective weights\n",
    "# Q, R, QN; and the controller horizon N) returns the cost\n",
    "# to go of the finite-horizon LQR using direct shooting\n",
    "def get_J_star_N_shooting(A, B, Q, R, QN, N, x0):\n",
    "    \n",
    "    # condensed dynamic equations\n",
    "    Ab = get_Ab(A, N)\n",
    "    Bb = get_Bb(A, B, N)\n",
    "    \n",
    "    # weight block matrices\n",
    "    Qb = sp.linalg.block_diag(*([Q] * (N - 1) + [QN]))\n",
    "    Rb = sp.linalg.block_diag(*([R] * N))\n",
    "    Hb = Rb + Bb.T.dot(Qb).dot(Bb)\n",
    "    \n",
    "    # solve for the optimal ub\n",
    "    ub = np.linalg.solve(Hb, - Bb.T.dot(Qb).dot(Ab).dot(x0))\n",
    "    \n",
    "    # plug optimal Ub in J to get the cost to go\n",
    "    J = x0.dot(Q + Ab.T.dot(Qb).dot(Ab)).dot(x0) + \\\n",
    "        2 * x0.dot(Ab.T).dot(Qb).dot(Bb).dot(ub) + \\\n",
    "        ub.dot(Hb).dot(ub)\n",
    "    \n",
    "    return J"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now plot the cost to go of the finite-horizon LQR problem as a function of the horizon $N$, and compare it with the infinite-horizon cost to go."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# function that plots the finite-horizon cost to go\n",
    "# as a function of the horizon length N\n",
    "# and compares it to the infinite-horizon solution\n",
    "# J_star_N is a list of floats: cost to go from N = 1 to N = N_max\n",
    "# J_star_inf is a float: infinite-horizon cost to go\n",
    "def plot_J_star(J_star_N, J_star_inf):\n",
    "    \n",
    "    # maximum length of the horizon\n",
    "    N_max = len(J_star_N)\n",
    "    \n",
    "    # infinite-horizon cost to go\n",
    "    plt.plot(range(1, N_max + 1), [J_star_inf] * N_max, label='Infinite-horizon LQR')\n",
    "    \n",
    "    # finite-horizon cost to go \n",
    "    plt.plot(range(1, N_max + 1), J_star_N, label='Finite-horizon LQR')\n",
    "    \n",
    "    # misc settings\n",
    "    plt.xlim(0, N_max)\n",
    "    plt.xlabel(r'Time horizon $N$')\n",
    "    plt.ylabel(r'Cost to go $J^*$')\n",
    "    plt.grid(True)\n",
    "    plt.legend()\n",
    "    \n",
    "# get cost to go as a function of N\n",
    "N_max = 45\n",
    "J_star_N_shooting = [get_J_star_N_shooting(A, B, Q, R, QN, N, x0) for N in range(1, N_max + 1)]\n",
    "\n",
    "# plot finite horizon vs infinite horizon\n",
    "plt.figure()\n",
    "plot_J_star(J_star_N_shooting, J_star_inf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The finite horizon solution already converges to the infinite-horizon one when $N$ is approximately equal to 5.\n",
    "But when $N$ gets too large, the direct shooting approach runs into troubles!\n",
    "\n",
    "To understand the reason for this behavior we analyze the [condition number](https://mathworld.wolfram.com/ConditionNumber.html) of $\\mathbb{H}$.\n",
    "For a matrix $A$ the condition number $\\text{cond}(A)$ is defined as the ratio between the largest and the smallest singular values.\n",
    "When solving a linear system $Ax = b$, the larger the condition number, the more sensitive is the solution $x = A^{-1} b$ to small perturbations of the vector $b$.\n",
    "Hence, if $\\text{cond}(A)$ is large, tiny numeric roundings in $b$ can be hugely amplified when solving for $x$.\n",
    "\n",
    "Another interesting interpretation is the following. The quantity\n",
    "$$\n",
    "\\log_2 (\\text{cond}(A))\n",
    "$$\n",
    "is an estimate of how many bits are lost in solving a linear system $Ax = b$.\n",
    "For example, if $\\text{cond}(A) = 2^{10}$ and $b$ is encoded in 64 bits, solving for $x$ we loose 10 of the 64 bits our machine is using!\n",
    "\n",
    "To get the optimal controls via direct shooting, we invert the matrix $\\mathbb{H}$.\n",
    "In the following cell, you are asked to analyze the condition number of $\\mathbb{H}$ as a function of the horizon length $N$.\n",
    "More specifically, you must set the variable `lost_bits_shooting` to be a list with `N_max = 45` elements.\n",
    "For `N` ranging from `1` to `N_max`, the `N`th entry in `lost_bits_shooting` must be logarithm in base 2 of the condition number of $\\mathbb{H}$ (`Hb` in the code) for a time horizon equal to `N`.\n",
    "\n",
    "P.s.: you might find the function [`numpy.linalg.cond`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.cond.html) useful here!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of bits lost in the inversion of Hb\n",
    "# as a function of the time horizon N\n",
    "lost_bits_shooting = [np.nan for N in range(1, N_max + 1)] # modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell we wrote a function to plot your results.\n",
    "Do you see anything alarming?\n",
    "Which kind of law does this curve follow: constant, logarithmic, linear, polynomial, or exponential?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the trend of the number of lost digits\n",
    "def plot_lost_bits(lost_bits):\n",
    "    \n",
    "    # plot number of bits in your machine\n",
    "    N_max = len(lost_bits)\n",
    "    plt.plot(\n",
    "        range(1, N_max + 1),\n",
    "        [np.log2(sys.maxsize) + 1] * N_max, # extra bit due to the sign\n",
    "        label='Bits of your machine'\n",
    "    )\n",
    "    \n",
    "    # plot number of bits lost inverting Hb\n",
    "    plt.plot(\n",
    "        range(1, N_max + 1),\n",
    "        lost_bits,\n",
    "        label=r'Bits lost inverting $\\mathbb{H}$'\n",
    "    )\n",
    "    \n",
    "    # misc settings\n",
    "    plt.xlim(0, N_max)\n",
    "    plt.xlabel(r'Time horizon $N$')\n",
    "    plt.ylabel(r'Number of bits')\n",
    "    plt.grid(True)\n",
    "    plt.legend()\n",
    "    \n",
    "# plot your results\n",
    "plt.figure()\n",
    "plot_lost_bits(lost_bits_shooting)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Direct Transcription\n",
    "\n",
    "Let's now give a try to the direct transcription approach.\n",
    "In these setting, our LQR problem is a quadratic optimization problem with equality constraints (i.e. the system dynamics).\n",
    "To solve it, we use the technique of the [Lagrange multipliers](https://en.wikipedia.org/wiki/Lagrange_multiplier).\n",
    "\n",
    "We define the Lagrange multipliers $\\lambda[0], \\ldots, \\lambda[N]$, and instead of enforcing the dynamics as constraints, we penalize it in the objective function.\n",
    "To this end, we define the Lagrangian function\n",
    "$$\n",
    "L =\n",
    "J +\n",
    "\\lambda^T[0] (\\mathbf{x}_0 - \\mathbf{x}[0]) +\n",
    "\\sum_{n=0}^{N-1}\n",
    "\\lambda^T[n+1](\\mathbf{A} \\mathbf{x}[n] + \\mathbf{B} \\mathbf{u}[n] - \\mathbf{x}[n+1]).\n",
    "$$\n",
    "\n",
    "The optimal solution of the constrained optimization problem, is now recovered by minimizing the unconstrained Lagrangian function $L$ with respect to the states $\\mathbf{x}[n]$, the controls $\\mathbf{u}[n]$, and also the multipliers $\\lambda[n]$.\n",
    "This can be done by simply setting the derivative of $L$ with respect to these variables to zero.\n",
    "\n",
    "Here are the resulting derivatives directly in matrix form.\n",
    "If your eyes are trained in taking derivatives, you might quickly double-check this, if not, don't feel the need to derive this again...\n",
    "\n",
    "<!--\n",
    "At the optimum, the gradient of $L$ with respect to $\\mathbf{x}[n]$, $\\mathbf{u}[n]$, and $\\lambda[n]$ must vanish.\n",
    "We obtain the optimality conditions:\n",
    "- $\\nabla_{\\lambda[0]} L = \\mathbf{x}_0 - \\mathbf{x}[0]$,\n",
    "- $\\nabla_{\\lambda[n]} L = \\mathbf{A} \\mathbf{x}[n] + \\mathbf{B} \\mathbf{u}[n] - \\mathbf{x}[n+1] = 0$ for $n = 1, \\ldots, N$,\n",
    "- $\\nabla_{\\mathbf{x}[n]} L = 2 \\mathbf{Q} \\mathbf{x}[n] - \\lambda[n] + \\mathbf{A}^T \\lambda[n+1] = 0$ for $n = 0, \\ldots, N-1$,\n",
    "- $\\nabla_{\\mathbf{x}[N]} L = 2 \\mathbf{Q} \\mathbf{x}[N] - \\lambda[N] = 0$,\n",
    "- $\\nabla_{\\mathbf{u}[n]} L = 2 \\mathbf{R} \\mathbf{u}[n] + \\mathbf{B}^T \\lambda[n+1] = 0$ for $n = 0, \\ldots, N-1$.\n",
    "-->\n",
    "\n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "\\nabla_{\\lambda[0]} L \\\\\n",
    "\\nabla_{\\mathbf{x}[0]} L \\\\\n",
    "\\nabla_{\\mathbf{u}[0]} L \\\\\n",
    "\\nabla_{\\lambda[1]} L \\\\\n",
    "\\nabla_{\\mathbf{x}[1]} L \\\\\n",
    "\\nabla_{\\mathbf{u}[1]} L \\\\\n",
    "\\vdots \\\\\n",
    "\\nabla_{\\mathbf{u}[N-1]} L \\\\\n",
    "\\nabla_{\\lambda[N]} L \\\\\n",
    "\\nabla_{\\mathbf{x}[N]} L \\\\\n",
    "\\end{bmatrix}\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "0 & -\\mathbf{I} & 0 & 0 & 0 & 0 & 0 & \\cdots & 0 & 0 & 0 & 0 \\\\\n",
    "-\\mathbf{I} & 2 \\mathbf{Q} & 0 & \\mathbf{A}^T & 0 & 0 & 0 & \\cdots & 0 & 0 & 0 & 0 \\\\\n",
    "0 & 0 & 2 \\mathbf{R} & \\mathbf{B}^T & 0 & 0 & 0 & \\cdots & 0 & 0 & 0 & 0 \\\\\n",
    "0 & \\mathbf{A} & \\mathbf{B} & 0 & -\\mathbf{I} & 0 & 0 & \\cdots & 0 & 0 & 0 & 0 \\\\\n",
    "0 & 0 & 0 & -\\mathbf{I} & 2 \\mathbf{Q} & 0 & \\mathbf{A}^T & \\cdots & 0 & 0 & 0 & 0 \\\\\n",
    "0 & 0 & 0 & 0 & 0 & 2 \\mathbf{R} & \\mathbf{B}^T & \\cdots & 0 & 0 & 0 & 0 \\\\\n",
    "\\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & \\ddots & \\vdots & \\vdots & \\vdots & \\vdots \\\\\n",
    "0 & 0 & 0 & 0 & 0 & 0 & 0 & \\cdots & 0 & 2 \\mathbf{R} & \\mathbf{B}^T & 0 \\\\\n",
    "0 & 0 & 0 & 0 & 0 & 0 & 0 & \\cdots & \\mathbf{A} & \\mathbf{B} & 0 & -\\mathbf{I} \\\\\n",
    "0 & 0 & 0 & 0 & 0 & 0 & 0 & \\cdots & 0 & 0 & -\\mathbf{I} & 2 \\mathbf{Q}_N \\\\\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "\\lambda[0] \\\\\n",
    "\\mathbf{x}[0] \\\\\n",
    "\\mathbf{u}[0] \\\\\n",
    "\\lambda[1] \\\\\n",
    "\\mathbf{x}[1] \\\\\n",
    "\\mathbf{u}[1] \\\\\n",
    "\\lambda[2] \\\\\n",
    "\\vdots \\\\\n",
    "\\mathbf{x}[N-1] \\\\\n",
    "\\mathbf{u}[N-1] \\\\\n",
    "\\lambda[N] \\\\\n",
    "\\mathbf{x}[N] \\\\\n",
    "\\end{bmatrix}\n",
    "+\n",
    "\\begin{bmatrix}\n",
    "\\mathbf{x}_0 \\\\\n",
    "0 \\\\\n",
    "0 \\\\\n",
    "0 \\\\\n",
    "0 \\\\\n",
    "0 \\\\\n",
    "\\vdots \\\\\n",
    "0 \\\\\n",
    "0 \\\\\n",
    "0 \\\\\n",
    "\\end{bmatrix}\n",
    "= 0.\n",
    "$$\n",
    "\n",
    "In the following, we will compactly refer to this big linear system as $\\mathbb{M} \\mathbb{z} = \\mathbb{x}_0$.\n",
    "Where $\\mathbb{M}$ is the big matrix in the center, $\\mathbb{z}$ are the states, the controls, and the multipliers, and $\\mathbb{x}_0$ is the (negative) initial conditions concatenated with all the zeros."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**On the sparsity pattern of the matrix $\\mathbb{M}$**\n",
    "\n",
    "Take some seconds to analyze the structure of the matrix $\\mathbb{M}$.\n",
    "\n",
    "First note that the entries of this matrix are nonzero only nearby the diagonal (such a matrix is called \"[band sparse matrix](https://en.wikipedia.org/wiki/Band_matrix)\").\n",
    "When a matrix is [sparse](https://en.wikipedia.org/wiki/Sparse_matrix), it's always a good idea to take advantage of that: knowing ahead of time where the nonzero entries of our matrix are can save a lot of redundant computations (mostly multiplications by zero).\n",
    "In the following code, we will exploit this structure in the solution of the linear system $\\mathbb{M} \\mathbb{z} = \\mathbb{x}_0$ by using [`scipy` sparse matrix package](https://docs.scipy.org/doc/scipy/reference/sparse.html).\n",
    "\n",
    "In the next cell, you are asked to code the building block of the function we use to construct $\\mathbb{M}$.\n",
    "For a generic time step $t$, consider the three equations: $\\nabla_{\\mathbf{x}[t]} L = 0$, $\\nabla_{\\mathbf{u}[t]} L = 0$, and $\\nabla_{\\lambda[t+1]} L = 0$.\n",
    "And note that these equations only depend on five variables: $\\lambda[t]$, $\\mathbf{x}[t]$, $\\mathbf{u}[t]$, $\\lambda[t+1]$, and $\\mathbf{x}[t+1]$.\n",
    "(To convince yourself, start by considering $t=0$.)\n",
    "\n",
    "These three equations and five variables determine a precise block of the matrix $\\mathbb{M}$ which is repeated $N-1$ times along the diagonal.\n",
    "Write a function named `diagonal_block_Mb` which constructs this block taking as inputs:\n",
    "- the linear system matrices $A$ and $B$,\n",
    "- the running cost weights $Q$ and $R$,\n",
    "\n",
    "To do this, use the `scipy` block matrix function [`scipy.sparse.bmat`](https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.sparse.bmat.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that given the dynamical system matrices\n",
    "# A and B, and the weights Q and R, returns the\n",
    "# reapeated block on the diagonal of Mb\n",
    "def Mb_diagonal_block(A, B, Q, R):\n",
    "    \n",
    "    # number of states and inputs of the linear dynamics\n",
    "    p, q = B.shape\n",
    "    \n",
    "    # block matrix\n",
    "    Ip = sp.sparse.eye(p) # modify here\n",
    "    Iq = sp.sparse.eye(q) # modify here\n",
    "    my_block = sp.sparse.bmat([ # modify here\n",
    "        [  Ip, None, None,   Ip, None], # modify here\n",
    "        [None, None,   Iq, None, None], # modify here\n",
    "        [None,   Ip, None, None,   Ip] # modify here\n",
    "        ])\n",
    "    \n",
    "    return my_block"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we fill the matrix $\\mathbb{M}$ with the code you wrote.\n",
    "We already took care of the first row (initial conditions) and the final row (terminal cost).\n",
    "For the rest we just copy and paste the matrix you wrote along the diagonal of $\\mathbb{M}$.\n",
    "\n",
    "Here are two efficient functions to derive $\\mathbb{M}$ and $\\mathbb{x}_0$ in [sparse form](https://en.wikipedia.org/wiki/Sparse_matrix)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# efficient computation of the matrix Mb\n",
    "# uses scipy sparse matrices:\n",
    "# https://docs.scipy.org/doc/scipy/reference/sparse.html\n",
    "def get_Mb(A, B, Q, R, QN, N):\n",
    "    \n",
    "    # number of states and inputs\n",
    "    p, q = B.shape\n",
    "    \n",
    "    # number of optimization variables in zb\n",
    "    n_zb = 2 * (N + 1) * p + N * q\n",
    "    \n",
    "    # initialize with empty sparse matrix\n",
    "    # (lil stands for list of lists)\n",
    "    Mb = sp.sparse.lil_matrix((n_zb, n_zb))\n",
    "    \n",
    "    # first block-row (initial conditions)\n",
    "    Mb[:p, p:2*p] = - sp.sparse.eye(p)\n",
    "    \n",
    "    # use the function you wrote to get the diagonal block\n",
    "    Mb_block = Mb_diagonal_block(A, B, Q, R)\n",
    "    \n",
    "    # initialize indices for the insertion of the blocks \n",
    "    rows = slice(p, Mb_block.shape[0] + p)\n",
    "    cols = slice(0, Mb_block.shape[1])\n",
    "    step = 2 * p + q\n",
    "    \n",
    "    # loop overt the diagonal and paste the Mb_block\n",
    "    # in the appropriate spots\n",
    "    for n in range(N):\n",
    "        Mb[rows, cols] = Mb_block\n",
    "        rows = slice(rows.start + step, rows.stop + step)\n",
    "        cols = slice(cols.start + step, cols.stop + step)\n",
    "\n",
    "    # last block-row (terminal state)\n",
    "    Mb[-p:, -2*p:-p] = - sp.sparse.eye(p)\n",
    "    Mb[-p:, -p:] = 2 * QN\n",
    "    \n",
    "    # convert matrix to scipy sparse csc matrix\n",
    "    # (csc stands for compressed sparse column)\n",
    "    return Mb.tocsc()\n",
    "\n",
    "# efficient computation of the vector xb0\n",
    "# p is the number of system states and\n",
    "# p is the number of system inputs\n",
    "def get_xb0(x0, p, q, N):\n",
    "    \n",
    "    # indices of the nonzero elements in xb0\n",
    "    rows = range(p)\n",
    "    cols = [0] * p\n",
    "    sparsity = (rows, cols)\n",
    "    \n",
    "    # overall shape of xb0\n",
    "    # (total number of rows, number of cols = 1)\n",
    "    n_zb = 2 * (N + 1) * p + N * q\n",
    "    shape = (n_zb, 1)\n",
    "    \n",
    "    # xb0 in sparse form\n",
    "    xb0 = sp.sparse.csc_matrix((x0, sparsity), shape=shape)\n",
    "    \n",
    "    return xb0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a nice `matplotlib` function that allows you to \"spy\" inside a matrix and look at its structure (a.k.a. the sparsity pattern)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# arbitrary short time horizon\n",
    "N = 5\n",
    "\n",
    "# plot the sparsity pattern of Mb\n",
    "plt.figure()\n",
    "plt.spy(get_Mb(A, B, Q, R, QN, N))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To simplify your upcoming work, we provided you:\n",
    "- Two functions to extract the states $\\mathbb{x}$ and the controls $\\mathbb{u}$ from the unknowns $\\mathbb{z}$ of the linear system $\\mathbb{M} \\mathbb{z} = \\mathbb{x}_0$.\n",
    "- A function to evaluate $J$ given $\\mathbb{x}$ and $\\mathbb{u}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that extracts the vector\n",
    "# xb = (x[0], ..., x[N]) from zb\n",
    "def extract_xb_from_zb(zb, p, q, N):\n",
    "    \n",
    "    # index of the first element of x[0] in the vector zb\n",
    "    start = p\n",
    "    \n",
    "    # distance between x[n] and x[n+1] in the vector zb\n",
    "    step = 2 * p + q\n",
    "    \n",
    "    # extract states\n",
    "    xb_indices = [start + step * n + i for n in range(N+1) for i in range(p)]\n",
    "    xb = zb[xb_indices]\n",
    "    \n",
    "    return xb\n",
    "\n",
    "# function that extracts the vector\n",
    "# ub = (u[0], ..., u[N-1]) from zb\n",
    "def extract_ub_from_zb(zb, p, q, N):\n",
    "    \n",
    "    # index of the first element of u[0] in the vector zb\n",
    "    start = 2 * p\n",
    "    \n",
    "    # distance between u[n] and u[n+1] in the vector zb\n",
    "    step = 2 * p + q\n",
    "    \n",
    "    # extract controls\n",
    "    ub_indices = [start + step * n + i for n in range(N) for i in range(q)]\n",
    "    ub = zb[ub_indices]\n",
    "\n",
    "    return ub\n",
    "\n",
    "# function that evaluates the objective J\n",
    "# given the block vectors xb and ub\n",
    "def evaluate_J(xb, ub, Q, R, QN, N):\n",
    "    \n",
    "    # block diagonal cost matrices\n",
    "    # (Qb has one extra block for the initial state here)\n",
    "    Qb = sp.sparse.block_diag(([Q] * N + [QN]))\n",
    "    Rb = sp.sparse.block_diag(([R] * N))\n",
    "    \n",
    "    # compute value function\n",
    "    J = xb.dot(Qb.dot(xb)) + ub.dot(Rb.dot(ub))\n",
    "    \n",
    "    return J"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's your turn again.\n",
    "For direct shooting, we have defined above the function `get_J_star_N_shooting(A, B, Q, R, QN, N, x0)`, which given the LQR data and the initial state, returns the cost to go for the finite-horizon problem using direct shooting.\n",
    "\n",
    "In the following cell you're asked to write the analogous function for the direct transcription method.\n",
    "The function must have the name `get_J_star_N_transcription` and its arguments must be `A, B, Q, R, QN, N, x0`.\n",
    "The steps in this function must be:\n",
    "- get the matrix $\\mathbb{M}$ and the vector $\\mathbb{x}_0$,\n",
    "- solve the linear system $\\mathbb{M} \\mathbb{z} = \\mathbb{x}_0$ for $\\mathbb{z}^*$,\n",
    "- extract $\\mathbb{x}^*$ and $\\mathbb{u}^*$ from $\\mathbb{z}^*$,\n",
    "- plug $\\mathbb{x}^*$ and $\\mathbb{u}^*$ in the expression for $J$.\n",
    "\n",
    "To solve the sparse linear system $\\mathbb{M} \\mathbb{z} = \\mathbb{x}_0$, use the `scipy` function [`scipy.sparse.linalg.spsolve`](https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.sparse.linalg.spsolve.html), which is already imported as `spsolve`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that given the system state x0 (together with:\n",
    "# the dynamical system matrices A, B; the objective weights\n",
    "# Q, R, QN; and the controller horizon N) returns the cost\n",
    "# to go of the finite-horizon LQR using direct transcription\n",
    "def get_J_star_N_transcription(A, B, Q, R, QN, N, x0):\n",
    "    return np.nan # modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the same plot we had for direct shooting, for a longer maximum horizon `N_max = 100`. Does it look better than the direct-shooting one?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get cost to go as a function of N\n",
    "N_max = 100\n",
    "J_star_N_transcription = [get_J_star_N_transcription(A, B, Q, R, QN, N, x0) for N in range(1, N_max + 1)]\n",
    "\n",
    "# plot finite horizon vs infinite horizon\n",
    "plt.figure()\n",
    "plot_J_star(J_star_N_transcription, J_star_inf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's time to replicate the plot of the lost digits in the case of direct transcription.\n",
    "Now we look at the condition number of $\\mathbb{M}$.\n",
    "\n",
    "Define a list with name `lost_bits_transcription` which contains `N_max = 100` elements.\n",
    "For `N` ranging from `1` to `N_max`, the `N`th entry in `lost_bits_transcription` must be logarithm in base 2 of the condition number of $\\mathbb{M}$ (`Mb` in the code) for a time horizon equal to `N`.\n",
    "Note that, in order to use the function [`numpy.linalg.cond`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.cond.html), you first need to convert `Mb` to be a dense matrix.\n",
    "This can be done as `Mb_dense = Mb.todense()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of bits lost in the inversion of Mb\n",
    "# as a function of the time horizon N\n",
    "lost_bits_transcription = [np.nan for N in range(1, N_max + 1)] # modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Does this plot look better than the direct-shooting one?\n",
    "Which kind of law does this curve follow: constant, logarithmic, linear, polynomial, or exponential?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot your results\n",
    "plt.figure()\n",
    "plot_lost_bits(lost_bits_transcription)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamic Programming Recursion, a.k.a. Riccati Recursion\n",
    "\n",
    "Even if we did our best to exploit the sparsity in the linear system of the direct transcription method, this approach was still quite slow.\n",
    "It turns out that the particular structure of the matrix $\\mathbb{M}$ can be exploited even more and the linear system $\\mathbb{M} \\mathbb{z} = \\mathbb{x}_0$ can be solved much faster.\n",
    "\n",
    "The name of the technique to do this is \"Riccati recursion.\"\n",
    "You can find more details in [this paper](https://link.springer.com/article/10.1023/A:1021711402723) or in [these lecture notes](https://web.stanford.edu/class/ee363/notes/riccati-derivation.pdf).\n",
    "(The second should be an easier read.)\n",
    "The main idea is to invert the matrix $\\mathbb{M}$ one block per time, starting from the bottom right, and moving up; in a similar spirit of dynamics programming.\n",
    "\n",
    "This idea has had a big impact in robotics, allowing the solution of optimal control problems in real time.\n",
    "Researchers from ETH Zurich have done an amazing job at pushing techniques like this to the limit (see, for example, [this paper](https://arxiv.org/pdf/1710.04029.pdf)).\n",
    "Here is their quadruped robot ANYmal.\n",
    "Its control software takes great advantage of techniques like the ones we studied in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import IFrame\n",
    "IFrame(src='https://www.youtube.com/embed/m1-s8iOJaI4', width='640', height='360')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Riccati recursion idea can be applied very broadly; all the way to nonlinear constrained optimization.\n",
    "However, in our simple settings, it can be seen that it is equivalent to the dynamic programming recursion [which we have seen in class](http://underactuated.mit.edu/lqr.html#dt_riccati).\n",
    "Have a look at [these lecture notes](https://web.stanford.edu/class/ee363/notes/riccati-derivation.pdf) if you want to convince yourself about this equivalence.\n",
    "\n",
    "We initialize $n = N$ and the cost-to-go matrix $\\mathbf{S}[N] = \\mathbf{Q}_N$.\n",
    "Then we go backwards in time computing\n",
    "$$\n",
    "\\mathbf{S}[n-1] =\n",
    "\\mathbf{Q} +\n",
    "\\mathbf{A}^T\\mathbf{S}[n]\\mathbf{A} -\n",
    "\\mathbf{A}^T\\mathbf{S}[n]\\mathbf{B}\n",
    "(\\mathbf{R} + \\mathbf{B}^T\\mathbf{S}[n]\\mathbf{B})^{-1}\n",
    "\\mathbf{B}^T\\mathbf{S}[n]\\mathbf{A}\n",
    "$$\n",
    "until we get $\\mathbf{S}[0]$.\n",
    "\n",
    "This is the last piece of code you are asked to write.\n",
    "A function called `riccati_recursion` which takes as inputs `(A, B, Q, R, QN, N)` and returns the optimal cost-to-go matrix $\\mathbf{S}[0]$ for the finite-horizon LQR with $N$ steps.\n",
    "Do this by implementing the recursion above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# implementation of the Riccati recursion\n",
    "def riccati_recursion(A, B, Q, R, QN, N):\n",
    "    S = QN\n",
    "    for n in range(N):\n",
    "        S = S # modify here\n",
    "    return S"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the same old plot of finite-horizon LQR vs infinite-horizon LQR.\n",
    "If you have done things correctly, this should be super fast and numerically very stable.\n",
    "We push the maximum time horizon `N_max` all the way to `1000` here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# implementation of the Riccati recursion\n",
    "# function that given the system state x0 (together with:\n",
    "# the dynamical system matrices A, B; the objective weights\n",
    "# Q, R, QN; and the controller horizon N) returns the cost\n",
    "# to go of the finite-horizon LQR\n",
    "def get_J_star_N_riccati(A, B, Q, R, QN, N, x0):\n",
    "    S = riccati_recursion(A, B, Q, R, QN, N)\n",
    "    return x0.dot(S).dot(x0)\n",
    "\n",
    "# get cost to go as a function of N\n",
    "N_max = 1000\n",
    "J_star_N_riccati = [get_J_star_N_riccati(A, B, Q, R, QN, N, x0) for N in range(1, N_max + 1)]\n",
    "\n",
    "# plot finite horizon vs infinite horizon\n",
    "plt.figure()\n",
    "plot_J_star(J_star_N_riccati, J_star_inf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Take-Home Messages\n",
    "\n",
    "Summarizing the main lessons of the notebook:\n",
    "- Linear algebra can be tedious, but is a fundamental issue in trajectory optimization.\n",
    "- Direct transcription leads to bigger optimization problems than direct shooting, but these are better conditioned and can be solved very fast using specialized techniques (Riccati recursion)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autograding\n",
    "You can check your work by running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underactuated.exercises.trajopt.shooting_vs_transcription.test_shooting_vs_transcription import TestShootingVsTranscription\n",
    "from underactuated.exercises.grader import Grader\n",
    "Grader.grade_output([TestShootingVsTranscription], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}